#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
增强版日志系统
支持动态倒计时、进度条、美观展示
"""

import time
import sys
import threading
from datetime import datetime, timedelta
from typing import Optional, Dict, Any
import numpy as np

try:
    from tqdm import tqdm
    TQDM_AVAILABLE = True
except ImportError:
    TQDM_AVAILABLE = False

class Colors:
    """ANSI颜色代码"""
    # 基础颜色
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    
    # 亮色
    BRIGHT_RED = '\033[1;91m'
    BRIGHT_GREEN = '\033[1;92m'
    BRIGHT_YELLOW = '\033[1;93m'
    BRIGHT_BLUE = '\033[1;94m'
    BRIGHT_MAGENTA = '\033[1;95m'
    BRIGHT_CYAN = '\033[1;96m'
    BRIGHT_WHITE = '\033[1;97m'
    
    # 背景色
    BG_RED = '\033[101m'
    BG_GREEN = '\033[102m'
    BG_YELLOW = '\033[103m'
    BG_BLUE = '\033[104m'
    
    # 样式
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    BLINK = '\033[5m'
    
    # 重置
    RESET = '\033[0m'
    CLEAR_LINE = '\033[2K\r'

class EnhancedLogger:
    """增强版日志记录器"""
    
    def __init__(self, name: str = "ProtoSeg"):
        self.name = name
        self.start_time = None
        self.current_iter = 0
        self.total_iters = 0
        self.current_epoch = 0
        self.total_epochs = 0
        self.metrics_history = []
        self._progress_thread = None
        self._stop_progress = False
        
    def start_training(self, total_iters: int, total_epochs: int = 1):
        """开始训练"""
        self.start_time = time.time()
        self.total_iters = total_iters
        self.total_epochs = total_epochs
        self.current_iter = 0
        self.current_epoch = 0
        
        self._print_header()
        self._start_progress_display()
        
    def _print_header(self):
        """打印训练头部信息"""
        print(f"\n{Colors.BRIGHT_CYAN}{'='*80}{Colors.RESET}")
        print(f"{Colors.BRIGHT_WHITE}🚀 {self.name} 训练开始{Colors.RESET}")
        print(f"{Colors.CYAN}📅 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.RESET}")
        print(f"{Colors.CYAN}🔄 总迭代数: {self.total_iters:,}{Colors.RESET}")
        print(f"{Colors.CYAN}📊 总轮数: {self.total_epochs}{Colors.RESET}")
        print(f"{Colors.BRIGHT_CYAN}{'='*80}{Colors.RESET}\n")
        
    def _start_progress_display(self):
        """启动进度显示线程"""
        if self._progress_thread is None or not self._progress_thread.is_alive():
            self._stop_progress = False
            self._progress_thread = threading.Thread(target=self._progress_loop, daemon=True)
            self._progress_thread.start()
            
    def _progress_loop(self):
        """进度显示循环"""
        while not self._stop_progress and self.current_iter < self.total_iters:
            self._update_progress_display()
            time.sleep(1)  # 每秒更新一次
            
    def _update_progress_display(self):
        """更新进度显示"""
        if self.current_iter >= self.total_iters:
            return
            
        # 计算进度
        progress = self.current_iter / self.total_iters
        elapsed_time = time.time() - self.start_time if self.start_time else 0
        
        # 计算剩余时间
        if self.current_iter > 0:
            eta = elapsed_time * (self.total_iters - self.current_iter) / self.current_iter
            eta_str = self._format_time(eta)
        else:
            eta_str = "计算中..."
            
        # 计算速度
        speed = self.current_iter / elapsed_time if elapsed_time > 0 else 0
        
        # 创建进度条
        bar_length = 40
        filled_length = int(bar_length * progress)
        bar = '█' * filled_length + '░' * (bar_length - filled_length)
        
        # 格式化输出
        progress_str = f"{progress*100:6.2f}%"
        iter_str = f"{self.current_iter:,}/{self.total_iters:,}"
        speed_str = f"{speed:.2f} it/s"
        elapsed_str = self._format_time(elapsed_time)
        
        # 构建显示字符串
        display = (
            f"{Colors.CLEAR_LINE}"
            f"{Colors.BRIGHT_BLUE}🔄 训练进度{Colors.RESET} "
            f"{Colors.BRIGHT_GREEN}[{bar}]{Colors.RESET} "
            f"{Colors.BRIGHT_WHITE}{progress_str}{Colors.RESET} "
            f"{Colors.CYAN}({iter_str}){Colors.RESET} "
            f"{Colors.YELLOW}⏱️  {elapsed_str}{Colors.RESET} "
            f"{Colors.MAGENTA}⏳ {eta_str}{Colors.RESET} "
            f"{Colors.BRIGHT_CYAN}🚀 {speed_str}{Colors.RESET}"
        )
        
        # 使用\r来覆盖当前行，而不是换行
        sys.stdout.write(f"\r{display}")
        sys.stdout.flush()
        
    def _format_time(self, seconds: float) -> str:
        """格式化时间显示"""
        if seconds < 60:
            return f"{seconds:.0f}s"
        elif seconds < 3600:
            return f"{seconds//60:.0f}m {seconds%60:.0f}s"
        else:
            hours = seconds // 3600
            minutes = (seconds % 3600) // 60
            return f"{hours:.0f}h {minutes:.0f}m"
            
    def update_iteration(self, iter_num: int, metrics: Optional[Dict[str, Any]] = None):
        """更新迭代信息"""
        self.current_iter = iter_num
        if metrics:
            self.metrics_history.append({
                'iter': iter_num,
                'time': time.time(),
                'metrics': metrics.copy()
            })
            
    def log_metrics(self, metrics: Dict[str, Any], prefix: str = ""):
        """记录指标"""
        if not metrics:
            return
            
        # 清空进度条并换行
        sys.stdout.write(f"{Colors.CLEAR_LINE}\n")
        
        # 构建指标字符串
        metric_strs = []
        for key, value in metrics.items():
            if isinstance(value, float):
                if 'loss' in key.lower():
                    color = Colors.RED if value > 1.0 else Colors.YELLOW if value > 0.5 else Colors.GREEN
                    metric_strs.append(f"{color}{key}: {value:.4f}{Colors.RESET}")
                elif 'acc' in key.lower() or 'iou' in key.lower() or 'miou' in key.lower():
                    color = Colors.GREEN if value > 0.8 else Colors.YELLOW if value > 0.5 else Colors.RED
                    metric_strs.append(f"{color}{key}: {value:.4f}{Colors.RESET}")
                else:
                    metric_strs.append(f"{Colors.CYAN}{key}: {value:.4f}{Colors.RESET}")
            else:
                metric_strs.append(f"{Colors.WHITE}{key}: {value}{Colors.RESET}")
                
        # 打印指标
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"{Colors.BRIGHT_BLUE}[{timestamp}]{Colors.RESET} {prefix}{' '.join(metric_strs)}")
        
        # 重新启动进度显示
        if not self._stop_progress:
            self._start_progress_display()
            
    def log_validation(self, metrics: Dict[str, Any]):
        """记录验证结果"""
        # 清空进度条
        sys.stdout.write(f"{Colors.CLEAR_LINE}")
        
        print(f"\n{Colors.BRIGHT_MAGENTA}📊 验证结果{Colors.RESET}")
        print(f"{Colors.MAGENTA}{'─'*50}{Colors.RESET}")
        
        for key, value in metrics.items():
            if isinstance(value, float):
                if 'loss' in key.lower():
                    color = Colors.RED if value > 1.0 else Colors.YELLOW if value > 0.5 else Colors.GREEN
                elif 'acc' in key.lower() or 'iou' in key.lower() or 'miou' in key.lower():
                    color = Colors.GREEN if value > 0.8 else Colors.YELLOW if value > 0.5 else Colors.RED
                else:
                    color = Colors.CYAN
                print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {color}{value:.4f}{Colors.RESET}")
            else:
                print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {Colors.CYAN}{value}{Colors.RESET}")
                
        print(f"{Colors.MAGENTA}{'─'*50}{Colors.RESET}\n")
        
        # 重新启动进度显示
        if not self._stop_progress:
            self._start_progress_display()
            
    def log_epoch_summary(self, epoch: int, train_metrics: Dict[str, Any], val_metrics: Optional[Dict[str, Any]] = None):
        """记录轮次总结"""
        self.current_epoch = epoch
        
        # 清空进度条
        sys.stdout.write(f"{Colors.CLEAR_LINE}")
        
        print(f"\n{Colors.BRIGHT_GREEN}📈 第 {epoch} 轮总结{Colors.RESET}")
        print(f"{Colors.GREEN}{'═'*60}{Colors.RESET}")
        
        # 训练指标
        print(f"{Colors.BRIGHT_YELLOW}🏋️  训练指标:{Colors.RESET}")
        for key, value in train_metrics.items():
            if isinstance(value, float):
                color = Colors.GREEN if 'acc' in key.lower() or 'iou' in key.lower() else Colors.YELLOW
                print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {color}{value:.4f}{Colors.RESET}")
            else:
                print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {Colors.CYAN}{value}{Colors.RESET}")
                
        # 验证指标
        if val_metrics:
            print(f"\n{Colors.BRIGHT_BLUE}🔍 验证指标:{Colors.RESET}")
            for key, value in val_metrics.items():
                if isinstance(value, float):
                    color = Colors.GREEN if 'acc' in key.lower() or 'iou' in key.lower() else Colors.BLUE
                    print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {color}{value:.4f}{Colors.RESET}")
                else:
                    print(f"  {Colors.WHITE}{key:15s}{Colors.RESET}: {Colors.CYAN}{value}{Colors.RESET}")
                    
        print(f"{Colors.GREEN}{'═'*60}{Colors.RESET}\n")
        
    def finish_training(self):
        """完成训练"""
        self._stop_progress = True
        
        # 清空进度条
        sys.stdout.write(f"{Colors.CLEAR_LINE}")
        
        total_time = time.time() - self.start_time if self.start_time else 0
        
        print(f"\n{Colors.BRIGHT_GREEN}🎉 训练完成!{Colors.RESET}")
        print(f"{Colors.GREEN}{'='*60}{Colors.RESET}")
        print(f"{Colors.CYAN}📅 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.RESET}")
        print(f"{Colors.CYAN}⏱️  总用时: {self._format_time(total_time)}{Colors.RESET}")
        print(f"{Colors.CYAN}🔄 总迭代: {self.current_iter:,}{Colors.RESET}")
        print(f"{Colors.CYAN}📊 平均速度: {self.current_iter/total_time:.2f} it/s{Colors.RESET}")
        print(f"{Colors.GREEN}{'='*60}{Colors.RESET}\n")
        
    def log_info(self, message: str, color: str = Colors.WHITE):
        """记录信息"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"{Colors.BRIGHT_BLUE}[{timestamp}]{Colors.RESET} {color}{message}{Colors.RESET}")
        
    def log_warning(self, message: str):
        """记录警告"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"{Colors.BRIGHT_YELLOW}[{timestamp}] ⚠️  {message}{Colors.RESET}")
        
    def log_error(self, message: str):
        """记录错误"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"{Colors.BRIGHT_RED}[{timestamp}] ❌ {message}{Colors.RESET}")
        
    def log_success(self, message: str):
        """记录成功"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"{Colors.BRIGHT_GREEN}[{timestamp}] ✅ {message}{Colors.RESET}")


# 全局日志实例
enhanced_logger = EnhancedLogger("ProtoSeg")
